#include "stdafx.h"
#include "filesys/process.hpp"
#include "filesys/environment.hpp"
#include "utils/numberformatter.hpp"
using namespace JHCPP::utils;
#include "filesys/pipe.hpp"
#include "stdlib/bugcheck.hpp"
#include "stdlib/exception.hpp"
using namespace JHCPP::stdlib::jhException;

namespace JHCPP
{
	namespace filesys
	{
		std::vector<char> getEnvironmentVariablesBuffer(const CProcess::Env& env)
		{   
			std::vector<char> envbuf;
			std::size_t pos = 0; 

			for (CProcess::Env::const_iterator it = env.begin(); it != env.end(); ++it)
			{
				std::size_t envlen = it->first.length() + it->second.length() + 1;

				envbuf.resize(pos + envlen + 1);
				std::copy(it->first.begin(), it->first.end(), &envbuf[pos]);
				pos += it->first.length();
				envbuf[pos] = '=';
				++pos;
				std::copy(it->second.begin(), it->second.end(), &envbuf[pos]);
				pos += it->second.length();

				envbuf[pos] = '\0';
				++pos;
			}

			envbuf.resize(pos + 1);
			envbuf[pos] = '\0';

			return envbuf;
		}

		void setEnvironmentVariables(const CProcess::Env& env)
		{
			for (CProcess::Env::const_iterator it = env.begin(); it != env.end(); ++it)
			{
				CEnvironment::set(it->first, it->second);
			}
		}

	#if defined(JH_OS_WINDOWS)
		//
		// CProcessHandleImpl
		//
		CProcessHandleImpl::CProcessHandleImpl(HANDLE hProcess, UInt32 pid) : m_hProcess(hProcess), m_pid(pid)
		{
		}

		CProcessHandleImpl::~CProcessHandleImpl()
		{
			::CloseHandle(m_hProcess);
		}

		UInt32 CProcessHandleImpl::id() const
		{
			return m_pid;
		}

		HANDLE CProcessHandleImpl::process() const
		{
			return m_hProcess;
		}

		int CProcessHandleImpl::wait() const
		{
			DWORD rc = ::WaitForSingleObject(m_hProcess, INFINITE);
			if (rc != WAIT_OBJECT_0)
				throw SystemException("Wait failed for process", CNumberFormatter::format(m_pid));

			DWORD exitCode;
			if (::GetExitCodeProcess(m_hProcess, &exitCode) == 0)
				throw SystemException("Cannot get exit code for process", CNumberFormatter::format(m_pid));

			return exitCode;
		}

		//
		// CProcessImpl
		//
		CProcessImpl::PIDImpl CProcessImpl::idImpl()
		{
			return ::GetCurrentProcessId(); 
		}

		void CProcessImpl::timesImpl(long& userTime, long& kernelTime)
		{
			FILETIME ftCreation;
			FILETIME ftExit;
			FILETIME ftKernel;
			FILETIME ftUser;

			if (::GetProcessTimes(::GetCurrentProcess(), &ftCreation, &ftExit, &ftKernel, &ftUser) != 0)
			{
				ULARGE_INTEGER time;
				time.LowPart  = ftKernel.dwLowDateTime;
				time.HighPart = ftKernel.dwHighDateTime;
				kernelTime    = long(time.QuadPart/10000000L);
				time.LowPart  = ftUser.dwLowDateTime;
				time.HighPart = ftUser.dwHighDateTime;
				userTime      = long(time.QuadPart/10000000L);
			}
			else
			{
				userTime = kernelTime = -1;
			}
		}

		CProcessHandleImpl* CProcessImpl::launchImpl(
			const std::string& command, 
			const ArgsImpl& args, 
			const std::string& initialDirectory, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe, 
			const EnvImpl& env)
		{
			std::string commandLine = command;
			for (ArgsImpl::const_iterator it = args.begin(); it != args.end(); ++it)
			{
				commandLine.append(" ");
				commandLine.append(*it);
			}		

			STARTUPINFOA startupInfo;
			::GetStartupInfo(&startupInfo); // take defaults from current process
			startupInfo.cb          = sizeof(STARTUPINFOA);
			startupInfo.lpReserved  = NULL;
			startupInfo.lpDesktop   = NULL;
			startupInfo.lpTitle     = NULL;
			startupInfo.dwFlags     = STARTF_FORCEOFFFEEDBACK;
			startupInfo.cbReserved2 = 0;
			startupInfo.lpReserved2 = NULL;

			HANDLE hProc = ::GetCurrentProcess();
			bool mustInheritHandles = false;
			if (inPipe)
			{
				::DuplicateHandle(hProc, inPipe->readHandle(), hProc, &startupInfo.hStdInput, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
				inPipe->close(CPipe::CLOSE_READ);
			}
			else if (::GetStdHandle(STD_INPUT_HANDLE))
			{
				::DuplicateHandle(hProc, ::GetStdHandle(STD_INPUT_HANDLE), hProc, &startupInfo.hStdInput, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
			}
			else 
			{
				startupInfo.hStdInput = 0;
			}
			// outPipe may be the same as errPipe, so we duplicate first and close later.
			if (outPipe)
			{
				::DuplicateHandle(hProc, outPipe->writeHandle(), hProc, &startupInfo.hStdOutput, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
			}
			else if (::GetStdHandle(STD_OUTPUT_HANDLE))
			{
				::DuplicateHandle(hProc, ::GetStdHandle(STD_OUTPUT_HANDLE), hProc, &startupInfo.hStdOutput, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
			}
			else 
			{
				startupInfo.hStdOutput = 0;
			}
			if (errPipe)
			{
				::DuplicateHandle(hProc, errPipe->writeHandle(), hProc, &startupInfo.hStdError, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
			}
			else if (::GetStdHandle(STD_ERROR_HANDLE))
			{
				::DuplicateHandle(hProc, ::GetStdHandle(STD_ERROR_HANDLE), hProc, &startupInfo.hStdError, 0, TRUE, DUPLICATE_SAME_ACCESS);
				mustInheritHandles = true;
			}
			else
			{
				startupInfo.hStdError = 0;
			}
			if (outPipe) outPipe->close(CPipe::CLOSE_WRITE);
			if (errPipe) errPipe->close(CPipe::CLOSE_WRITE);

			if (mustInheritHandles)
			{
				startupInfo.dwFlags |= STARTF_USESTDHANDLES;
			}

			const char* workingDirectory = initialDirectory.empty() ? 0 : initialDirectory.c_str();

			const char* pEnv = 0;
			std::vector<char> envChars;
			if (!env.empty())
			{
				envChars = getEnvironmentVariablesBuffer(env);
				pEnv = &envChars[0];
			}

			PROCESS_INFORMATION processInfo;
			DWORD creationFlags = ::GetConsoleWindow() ? 0 : CREATE_NO_WINDOW;
			BOOL rc = ::CreateProcess(
				NULL, 
				const_cast<char*>(commandLine.c_str()), 
				NULL, // processAttributes
				NULL, // threadAttributes
				mustInheritHandles,
				creationFlags,
				(LPVOID) pEnv, 
				workingDirectory, 
				&startupInfo, 
				&processInfo
				);
			if (startupInfo.hStdInput) ::CloseHandle(startupInfo.hStdInput);
			if (startupInfo.hStdOutput) ::CloseHandle(startupInfo.hStdOutput);
			if (startupInfo.hStdError) ::CloseHandle(startupInfo.hStdError);
			if (rc)
			{
				::CloseHandle(processInfo.hThread);
				return new CProcessHandleImpl(processInfo.hProcess, processInfo.dwProcessId);
			}
			else throw SystemException("Cannot launch process", command);
		}

		void CProcessImpl::killImpl(const CProcessHandleImpl& handle)
		{
			if (::TerminateProcess(handle.process(), 0) == 0)
			{
				::CloseHandle(handle.process());
				throw SystemException("cannot kill process");
			}
			::CloseHandle(handle.process());
		}

		void CProcessImpl::killImpl(PIDImpl pid)
		{
			HANDLE hProc = ::OpenProcess(PROCESS_TERMINATE, FALSE, pid);
			if (hProc)
			{
				if (::TerminateProcess(hProc, 0) == 0)
				{
					::CloseHandle(hProc);
					throw SystemException("cannot kill process");
				}
				::CloseHandle(hProc);
			}
			else
			{
				switch (GetLastError())
				{
				case ERROR_ACCESS_DENIED:
					throw NoPermissionException("cannot kill process");
				case ERROR_NOT_FOUND: 
					throw NotFoundException("cannot kill process");
				default:
					throw SystemException("cannot kill process");
				}
			}
		}

		void CProcessImpl::requestTerminationImpl(PIDImpl pid)
		{
			//NamedEvent ev(terminationEventName(pid));
			//ev.set();
		}

		std::string CProcessImpl::terminationEventName(PIDImpl pid)
		{
			std::string evName("POCOTRM");
			CNumberFormatter::appendHex(evName, pid, 8);
			return evName;
		}
	#elif defined(JH_OS_LINUX)
		//
		// CProcessHandleImpl
		//
		CProcessHandleImpl::CProcessHandleImpl(pid_t pid) : m_pid(pid)
		{
		}

		CProcessHandleImpl::~CProcessHandleImpl()
		{
		}

		pid_t CProcessHandleImpl::id() const
		{
			return m_pid;
		}

		int CProcessHandleImpl::wait() const
		{
			int status;
			int rc;
			do
			{
				rc = waitpid(m_pid, &status, 0);
			}
			while (rc < 0 && errno == EINTR);
			if (rc != m_pid)
				throw SystemException("Cannot wait for process", CNumberFormatter::format(_pid));
			return WEXITSTATUS(status);
		}

		//
		// CProcessImpl
		//
		CProcessImpl::PIDImpl CProcessImpl::idImpl()
		{
			return getpid();
		}

		void CProcessImpl::timesImpl(long& userTime, long& kernelTime)
		{
			struct rusage usage;
			getrusage(RUSAGE_SELF, &usage);
			userTime   = usage.ru_utime.tv_sec;
			kernelTime = usage.ru_stime.tv_sec;
		}

		CProcessHandleImpl* CProcessImpl::launchImpl(
			const std::string& command, 
			const ArgsImpl& args, 
			const std::string& initialDirectory, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe, 
			const EnvImpl& env)
		{
			return launchByForkExecImpl(command, args, initialDirectory, inPipe, outPipe, errPipe, env);
		}

		CProcessHandleImpl* CProcessImpl::launchByForkExecImpl(
			const std::string& command, 
			const ArgsImpl& args, 
			const std::string& initialDirectory, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe, 
			const EnvImpl& env)
		{
			int pid = fork();
			if (pid < 0)
			{
				throw SystemException("Cannot fork process for", command);		
			}
			else if (pid == 0)
			{
				if (!initialDirectory.empty())
				{
					if (chdir(initialDirectory.c_str()) != 0)
					{
						_exit(72);
					}
				}

				setEnvironmentVariables(env);

				// setup redirection
				if (inPipe)
				{
					dup2(inPipe->readHandle(), STDIN_FILENO);
					inPipe->close(CPipe::CLOSE_BOTH);
				}
				// outPipe and errPipe may be the same, so we dup first and close later
				if (outPipe) dup2(outPipe->writeHandle(), STDOUT_FILENO);
				if (errPipe) dup2(errPipe->writeHandle(), STDERR_FILENO);
				if (outPipe) outPipe->close(CPipe::CLOSE_BOTH);
				if (errPipe) errPipe->close(CPipe::CLOSE_BOTH);
				// close all open file descriptors other than stdin, stdout, stderr
				for (int i = 3; i < getdtablesize(); ++i)
					close(i);

				char** argv = new char*[args.size() + 2];
				int i = 0;
				argv[i++] = const_cast<char*>(command.c_str());
				for (ArgsImpl::const_iterator it = args.begin(); it != args.end(); ++it) 
					argv[i++] = const_cast<char*>(it->c_str());
				argv[i] = NULL;
				execvp(command.c_str(), argv);
				_exit(72);
			}

			if (inPipe)  inPipe->close(CPipe::CLOSE_READ);
			if (outPipe) outPipe->close(CPipe::CLOSE_WRITE);
			if (errPipe) errPipe->close(CPipe::CLOSE_WRITE);
			return new ProcessHandleImpl(pid);
		}

		void CProcessImpl::killImpl(const CProcessHandleImpl& handle)
		{
			killImpl(handle.id());
		}

		void CProcessImpl::killImpl(PIDImpl pid)
		{
			if (kill(pid, SIGKILL) != 0)
			{
				switch (errno)
				{
				case ESRCH:
					throw NotFoundException("cannot kill process");
				case EPERM:
					throw NoPermissionException("cannot kill process");
				default:
					throw SystemException("cannot kill process");
				}
			}
		}

		void CProcessImpl::requestTerminationImpl(PIDImpl pid)
		{
			if (kill(pid, SIGINT) != 0)
			{
				switch (errno)
				{
				case ESRCH:
					throw NotFoundException("cannot terminate process");
				case EPERM:
					throw NoPermissionException("cannot terminate process");
				default:
					throw SystemException("cannot terminate process");
				}
			}
		}
	#endif

		//
		// CProcessHandle
		//
		CProcessHandle::CProcessHandle(const CProcessHandle& handle) : m_pImpl(handle.m_pImpl)
		{
			m_pImpl->add();
		}

		CProcessHandle::~CProcessHandle()
		{
			m_pImpl->release();
		}

		CProcessHandle::CProcessHandle(CProcessHandleImpl* pImpl) : m_pImpl(pImpl)
		{
			jh_check_ptr (m_pImpl);
		}

		CProcessHandle& CProcessHandle::operator = (const CProcessHandle& handle)
		{
			if (&handle != this)
			{
				m_pImpl->release();
				m_pImpl = handle.m_pImpl;
				m_pImpl->add();
			}
			return *this;
		}

		CProcessHandle::PID CProcessHandle::id() const
		{
			return m_pImpl->id();
		}

		int CProcessHandle::wait() const
		{
			return m_pImpl->wait();
		}

		//
		// CProcess
		//
		CProcess::PID CProcess::id()
		{
			return CProcessImpl::idImpl();
		}

		void CProcess::times(long& userTime, long& kernelTime)
		{
			CProcessImpl::timesImpl(userTime, kernelTime);
		}

		CProcessHandle CProcess::launch(const std::string& command, const Args& args)
		{
			std::string initialDirectory;
			Env env;
			return CProcessHandle(launchImpl(command, args, initialDirectory, 0, 0, 0, env));
		}

		CProcessHandle CProcess::launch(
			const std::string& command, 
			const Args& args, 
			const std::string& initialDirectory)
		{
			Env env;
			return CProcessHandle(launchImpl(command, args, initialDirectory, 0, 0, 0, env));
		}

		CProcessHandle CProcess::launch(
			const std::string& command, 
			const Args& args, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe)
		{
			jh_assert (inPipe == 0 || (inPipe != outPipe && inPipe != errPipe));
			std::string initialDirectory;
			Env env;
			return CProcessHandle(launchImpl(command, args, initialDirectory, inPipe, outPipe, errPipe, env));
		}

		CProcessHandle CProcess::launch(
			const std::string& command, 
			const Args& args, 
			const std::string& initialDirectory, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe)
		{
			jh_assert (inPipe == 0 || (inPipe != outPipe && inPipe != errPipe));
			Env env;
			return CProcessHandle(launchImpl(command, args, initialDirectory, inPipe, outPipe, errPipe, env));
		}

		CProcessHandle CProcess::launch(
			const std::string& command, 
			const Args& args, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe, 
			const Env& env)
		{
			jh_assert (inPipe == 0 || (inPipe != outPipe && inPipe != errPipe));
			std::string initialDirectory;
			return CProcessHandle(launchImpl(command, args, initialDirectory, inPipe, outPipe, errPipe, env));
		}

		CProcessHandle CProcess::launch(
			const std::string& command,
			const Args& args,
			const std::string& initialDirectory, 
			CPipe* inPipe, 
			CPipe* outPipe, 
			CPipe* errPipe, 
			const Env& env)
		{
			jh_assert (inPipe == 0 || (inPipe != outPipe && inPipe != errPipe));
			return CProcessHandle(launchImpl(command, args, initialDirectory, inPipe, outPipe, errPipe, env));
		}

		int CProcess::wait(const CProcessHandle& handle)
		{
			return handle.wait();
		}

		void CProcess::kill(const CProcessHandle& handle)
		{
			killImpl(*handle.m_pImpl);
		}

		void CProcess::kill(PID pid)
		{
			killImpl(pid);
		}

		void CProcess::requestTermination(PID pid)
		{
			requestTerminationImpl(pid);
		}
	}//end of namespace filesys
}//end of namespace JHCPP